# -*- coding: utf-8 -*-
"""
Created on Tue Apr 27 13:53:27 2021

@author: zhiyinan
"""
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
# import Parameters_1D as par_RNN 
# from Polar1D_vectorized import *
# import van_Genuchten1D_vectorized as vg 
import Parameters1D_casadi as par 
from Model1D_casadi_crop import *
import van_Genuchten1D_casadi as vg 

from scipy import integrate
from sklearn.preprocessing import MinMaxScaler
from casadi import *
import matplotlib.font_manager as font_manager
import time
import pandas as pd
from sklearn.linear_model import Ridge
import random
# %% Define the model 
# NNmodel = tf.keras.models.load_model('past5_future5_10_5sigmoid_LSTM_1hr')
# NNmodel = tf.keras.models.load_model('past5_future1_5sigmoid_3t_LSTM_1hr')

# weightLSTM_1 = NNmodel.layers[0].get_weights()
# warr1, uarr1, barr1 = weightLSTM_1
# weightLSTM_2 = NNmodel.layers[1].get_weights()
# warr2,uarr2, barr2 = weightLSTM_2

# weightD = NNmodel.layers[2].get_weights()
# kernel3, bias3 = weightD

def sigmoid(x):
    return 1/(1+np.exp(-x))

def LSTMlayer_tanh(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    _c = np.tanh(s_t[:,2*hunit:3*hunit])
    # _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    h_t = o*np.tanh(c_t)
    # h_t = o*sigmoid(c_t)
    return h_t, c_t

def LSTMlayer_sig(warr, uarr, barr, x_t,h_tm1,c_tm1):
    
    '''
    c_tm1 = np.array([0,0]).reshape(1,2)
    h_tm1 = np.array([0,0]).reshape(1,2)
    x_t   = np.array([1]).reshape(1,1)
    
    warr.shape = (nfeature,hunits*4)
    uarr.shape = (hunits,hunits*4)
    barr.shape = (hunits*4,)
    
    '''
    
    s_t = (mtimes(x_t,warr) + mtimes(h_tm1,uarr) + barr.reshape((1, barr.shape[0])))
    hunit = uarr.shape[0]
    i  = sigmoid(s_t[:,:hunit])
    f  = sigmoid(s_t[:,1*hunit:2*hunit])
    # _c = np.tanh(s_t[:,2*hunit:3*hunit])
    _c = sigmoid(s_t[:,2*hunit:3*hunit])
    o  = sigmoid(s_t[:,3*hunit:])
    c_t = i*_c + f*c_tm1
    # h_t = o*np.tanh(c_t)
    h_t = o*sigmoid(c_t)
    return h_t, c_t


#%% Define Model 1

model_1 = tf.keras.models.load_model("Modeling/LSTM_5s_5s_t_10nt_2")
# np.savetxt("LSTM_warr1",warr1)
# np.savetxt("LSTM_uarr1",uarr1)
# np.savetxt("LSTM_barr1",barr1)

# np.savetxt("LSTM_warr2",warr2)
# np.savetxt("LSTM_uarr2",uarr2)
# np.savetxt("LSTM_barr2",barr2)

# np.savetxt("LSTM_ker3", kernel3)
# np.savetxt("LSTM_bis3", bias3)
def M(x):
    weightLSTM_1 = model_1.layers[0].get_weights()
    warr1, uarr1, barr1 = weightLSTM_1

    weightLSTM_2 = model_1.layers[1].get_weights()
    warr2, uarr2, barr2 = weightLSTM_2
    
    weightD = model_1.layers[2].get_weights()
    kernel3, bias3 = weightD
    
    n1 = 5
    # Initialise the model with the hidden states
    h_t1 = SX.zeros((1, n1))
    # Long term memory
    C_t1 = SX.zeros((1, n1))
    for i in range(x.shape[0]):
        xi = x[i,:].reshape((1,2))
        L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
        h_t1 = vertcat(h_t1, L1[0])
        C_t1 = vertcat(C_t1, L1[1])
    
    n2 = 5
    # Initialise the model with the hidden states
    h_t2 = SX.zeros((1, n2))
    # Long term memory
    C_t2 = SX.zeros((1, n2))
    for i in range(h_t1.shape[0]):
        xi = h_t1[i,:].reshape((1,n1))
        L2 = LSTMlayer_sig(warr2, uarr2, barr2, xi, h_t2[i,:], C_t2[i,:])
        h_t2 = vertcat(h_t2, L2[0])
        C_t2 = vertcat(C_t2, L2[1])
    #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
    # Dense layer
    out2 = []
    for i in range(bias3.shape[0]):
        temp2 = mtimes(h_t2[-1,:], kernel3)[i]
    temp2 = bias3.reshape((1, bias3.shape[0])) + mtimes(h_t2[-1,:], kernel3)
    out2 = tanh(temp2)     
    
    return out2


# # %% Define Model 1
# model_1 = tf.keras.models.load_model("Modeling/p20_f1_10s_sig_2h_012_027_5_150_e8_n")

# def M1(x):
#     weightLSTM_1 = model_1.layers[0].get_weights()
#     warr1, uarr1, barr1 = weightLSTM_1

#     weightD = model_1.layers[1].get_weights()
#     kernel2, bias2 = weightD
    
#     n1 = 10
#     # Initialise the model with the hidden states
#     h_t1 = SX.zeros((1, n1))
#     # Long term memory
#     C_t1 = SX.zeros((1, n1))
#     for i in range(x.shape[0]):
#         xi = x[i,:].reshape((1,2))
#         L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
#         h_t1 = vertcat(h_t1, L1[0])
#         C_t1 = vertcat(C_t1, L1[1])
    
#     #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
#     # Dense layer
#     out2 = []
#     for i in range(bias2.shape[0]):
#         temp2 = mtimes(h_t1[-1,:], kernel2)[i]
#     temp2 = bias2.reshape((1, bias2.shape[0])) + mtimes(h_t1[-1,:], kernel2)
#     out2 = sigmoid(temp2)     
    
#     return out2
    
# # %% Define sub model 2
# model_2 = tf.keras.models.load_model('Modeling/p20_f1_10s_tanh_2h_021_032_8_90_e8_n')

# def M2(x):
#     weightLSTM_1 = model_2.layers[0].get_weights()
#     warr1, uarr1, barr1 = weightLSTM_1

#     weightD = model_2.layers[1].get_weights()
#     kernel2, bias2 = weightD
    
#     n1 = 10
#     # Initialise the model with the hidden states
#     h_t1 = SX.zeros((1, n1))
#     # Long term memory
#     C_t1 = SX.zeros((1, n1))
#     for i in range(x.shape[0]):
#         xi = x[i,:].reshape((1,2))
#         L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
#         h_t1 = vertcat(h_t1, L1[0])
#         C_t1 = vertcat(C_t1, L1[1])
    
#     #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
#     # Dense layer
#     out2 = []
#     for i in range(bias2.shape[0]):
#         temp2 = mtimes(h_t1[-1,:], kernel2)[i]
#     temp2 = bias2.reshape((1, bias2.shape[0])) + mtimes(h_t1[-1,:], kernel2)
#     out2 = tanh(temp2)     
    
#     return out2

# # %% Define sub model 3
# model_3 = tf.keras.models.load_model("Modeling/p20_f1_5s_3s_sig_2h_029_038_4_35_e7_n")

# # np.savetxt("Modeling/LSTM3_029_038_warr1.txt", warr1)
# # np.savetxt("Modeling/LSTM3_029_038_uarr1.txt", uarr1)
# # np.savetxt("Modeling/LSTM3_029_038_barr1.txt", barr1)

# # np.savetxt("Modeling/LSTM3_029_038_warr2.txt", warr2)
# # np.savetxt("Modeling/LSTM3_029_038_uarr2.txt", uarr2)
# # np.savetxt("Modeling/LSTM3_029_038_barr2.txt", barr2)

# # np.savetxt("Modeling/LSTM3_029_038_ker3.txt", kernel3)
# # np.savetxt("Modeling/LSTM3_029_038_bis3.txt", bias3)
# def M3(x):
#     weightLSTM_1 = model_3.layers[0].get_weights()
#     warr1, uarr1, barr1 = weightLSTM_1

#     weightLSTM_2 = model_3.layers[1].get_weights()
#     warr2, uarr2, barr2 = weightLSTM_2
    
#     weightD = model_3.layers[2].get_weights()
#     kernel3, bias3 = weightD
    
#     n1 = 5
#     # Initialise the model with the hidden states
#     h_t1 = SX.zeros((1, n1))
#     # Long term memory
#     C_t1 = SX.zeros((1, n1))
#     for i in range(x.shape[0]):
#         xi = x[i,:].reshape((1,2))
#         L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
#         h_t1 = vertcat(h_t1, L1[0])
#         C_t1 = vertcat(C_t1, L1[1])
    
#     n2 = 3
#     # Initialise the model with the hidden states
#     h_t2 = SX.zeros((1, n2))
#     # Long term memory
#     C_t2 = SX.zeros((1, n2))
#     for i in range(h_t1.shape[0]):
#         xi = h_t1[i,:].reshape((1,n1))
#         L2 = LSTMlayer_sig(warr2, uarr2, barr2, xi, h_t2[i,:], C_t2[i,:])
#         h_t2 = vertcat(h_t2, L2[0])
#         C_t2 = vertcat(C_t2, L2[1])
#     #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
#     # Dense layer
#     out2 = []
#     for i in range(bias3.shape[0]):
#         temp2 = mtimes(h_t2[-1,:], kernel3)[i]
#     temp2 = bias3.reshape((1, bias3.shape[0])) + mtimes(h_t2[-1,:], kernel3)
#     out2 = sigmoid(temp2)
    
#     return out2

# # %%
# model_a = tf.keras.models.load_model("Modeling/2nd_layer_4in_5s_s_sub_noise")
# # np.savetxt("NN_para/Ma_ker1.txt", kernel1)
# # np.savetxt("NN_para/Ma_bis1.txt", bias1)
# # np.savetxt("NN_para/Ma_ker2.txt", kernel2)
# # np.savetxt("NN_para/Ma_bis2.txt", bias2)

# def Ma(x):
#     weightD1 = model_a.layers[0].get_weights()
#     kernel1, bias1 = weightD1

#     weightD2 = model_a.layers[1].get_weights()
#     kernel2, bias2 = weightD2
    
#     # weightD3 = model_a.layers[2].get_weights()
#     # kernel3, bias3 = weightD3
    
#     # weightD4 = model_a.layers[3].get_weights()
#     # kernel4, bias4 = weightD4
    
#     # n1 = 3
#     # out1 = []
#     # for i in range(bias1.shape[0]):
#     #     temp1 = mtimes(x, kernel1)[i]
#     temp1 = bias1.reshape((1, bias1.shape[0])) + mtimes(x.reshape((1,x.shape[0])), kernel1)
#     out1 = sigmoid(temp1)   
    
#     temp2 = bias2.reshape((1, bias2.shape[0])) + mtimes(out1, kernel2)
#     out2 = sigmoid(temp2)     
    
#     # temp3 = bias3.reshape((1, bias3.shape[0])) + mtimes(out2, kernel3)
#     # out3 = tanh(temp3)  
    
#     # temp4 = bias4.reshape((1, bias4.shape[0])) + mtimes(out3, kernel4)
#     # out4 = sigmoid(temp4)  
    
#     return out2


# %% Define functions that defines the Zone MPC problem with given conditions 
Nu = 1
Ny = 1
Np = 1
# NN_in = vertcat(NN_in, horzcat(uc[0,:], yc[0,:]))

# def obj(u, y, yz, Q, R):
#     J = mtimes((y - yz).T, mtimes(Q, (y - yz))) + R*(1 - u)**2
#     return J

def obj(u, yl, yu, Ql, Qu, R):
    J = mtimes(yl.T, mtimes(Ql, yl)) + mtimes(yu.T, mtimes(Qu, yu)) + R*(1 - u)**2
    return J

# def zoneMPC(u, Precip, y, yz, ug, yg, lbu, ubu, lby, uby, lbZone, ubZone, 
#             Q, R, NN_past, y0, Nt, Delta, e_a, e_b):
def zoneMPC(u, Precip, y, yl, yu, ug, yg, lbu, ubu, lby, uby, lbZone, ubZone, 
              QL, QU, R, NN_past, y0, Nt, Delta, e_a, e_b):
    
    # Define the scaled variabes
    # uc = []
    yc = []
    # lbx = np.array([lbu,lby])
    # ubx = np.array([ubu,uby])
    # for i in range(int(Nt/Delta)):
    #     uci = []
    #     for j in range(Nu):
    #         uci = horzcat(uci, (u[i,j] - lbu)/(ubu - lbu))
    #         uc = vertcat(uc,uci)
    
    for i in range(int(Nt/Delta)+1):
        yci = []
        for j in range(Ny):    
            yci = horzcat(yci, (y[i,j] - lby)/(uby - lby)) # - 0.1)/(0.38 - 0.1))
            yc = vertcat(yc, yci)
    
    
    # for i in range(1,int(Nt/Delta)):
    #     NN_in = vertcat(NN_in, horzcat(u[i], yc[0]))
    # for i in range(future - int(Nt/Delta)):
    #     NN_in = vertcat(NN_in, horzcat(u[-1], yc[0])) 
    # model_sequential = model(NN_in)
    
    # Define the optimization problem
    W = []
    W0 = []
    W_lb = []
    W_ub = []
    f = 0
    g = []
    g_lb = []
    g_ub = []
    P = []
    
    W = vertcat(W,y[0,:])
    W_lb = vertcat(W_lb,y0)
    W_ub = vertcat(W_ub,y0)
    W0 = vertcat(W0,y0)
    
    # Predict the future outputs using the explictly expressed NN
    NN_in = NN_past
    
    for t in range(int(Nt/Delta)):
        # System inputs
        W = horzcat(W,u[t,:])
        W_lb = vertcat(W_lb,0)
        W_ub = vertcat(W_ub,1)
        W0 = vertcat(W0,ug)

        # Predicted system outputs
        W = horzcat(W,y[t+1,:])
        W_lb = vertcat(W_lb, lby)
        W_ub = vertcat(W_ub, uby)
        W0 = vertcat(W0,yg)
        
        # Zone slack y
        # W = horzcat(W,yz[t,:])
        # W_lb = vertcat(W_lb, lbZone[t])
        # W_ub = vertcat(W_ub, ubZone[t])
        # W0 = vertcat(W0, lbZone[t])
        
        # Parameter - Weather forecasting (scaled)
        # P = vertcat(P, (Precip[t] - lbu)/(ubu-lbu))
        P = vertcat(P, Precip[t])
        
        # Lower Slack
        W = horzcat(W,yl[t,:])
        W_lb = vertcat(W_lb, 0)
        W_ub = vertcat(W_ub, 0.2)
        W0 = vertcat(W0, lbZone[t])
        
        # Upper Slack
        W = horzcat(W,yu[t,:])
        W_lb = vertcat(W_lb, 0)
        W_ub = vertcat(W_ub, 0.2)
        W0 = vertcat(W0, ubZone[t])
        
        
        # System model constraints
        water = ((u[t]*(ubu-lbu) + lbu + Precip[t]) - lbu)/(ubu-lbu)
        # water = ((-(u[t]*(ubu-lbu) + lbu) + Precip[t]) + ubu)/(ubu-lbu)
        NN_in = vertcat(NN_in[1:,:], horzcat(water, yc[t]))
        # temp = []
        # temp = Ma(vertcat(M1(NN_in), M2(NN_in), M3(NN_in), yc[t]))*(uby - lby) +lby  #*(0.38 - 0.1) + 0.1  # , yc[t]
        temp = M(NN_in)*(uby - lby) +lby
        g = vertcat(g, temp*e_a + e_b - y[t+1,:])
        g_lb = vertcat(g_lb, 0)
        g_ub = vertcat(g_ub, 0)
        
        # Upper Zone constraints
        g = vertcat(g, y[t+1,:] + yl[t,:] - yu[t,:])
        g_lb = vertcat(g_lb, lbZone[t])
        g_ub = vertcat(g_ub, ubZone[t])
        
        # # Lower Zone constraints
        # g = vertcat(g, y[t+1,:] + yl[t,:])
        # g_lb = vertcat(g_lb, lbZone[t])
        # g_ub = vertcat(g_ub, 0.5)
        
        # Objective function
        f += obj(u[t,:], yl[t,:], yu[t,:], QL, QU, R)
        
    
    return W.T, P, W0, W_lb, W_ub, f, g, g_lb, g_ub 



# %% Error handling

def error_model(y_sim, y_pred):
    W = []
    W0 = []
    W_lb = []
    W_ub = []
    f = 0
    # g = []
    # g_lb = []
    # g_ub = []
    
    coef = SX.sym('coef', 2)
    W = vertcat(W, coef[0])
    W0 = vertcat(W0, 1)
    W_lb = vertcat(W_lb, 0.9)
    W_ub = vertcat(W_ub, 1.1)
    
    W = vertcat(W, coef[1])
    W0 = vertcat(W0, 0)
    W_lb = vertcat(W_lb, -0.05)
    W_ub = vertcat(W_ub, 0.05)
    n = y_pred.shape[0]
    for i in range(n):
        f += (y_sim[i] - (y_pred[i]*coef[0]+coef[1]))**2
    nlp = {'x':W,'f':f}
    S = nlpsol('S','ipopt',nlp)
    r = S(x0 = W0, lbx = W_lb, ubx = W_ub) 
    return r['x'][0], r['x'][1]

# %% 
random.seed(5)
np.random.seed(5)
past = 20
future = 1

y_index = 19

Delta = 2*60*60
Nt = 40*60*60
Nsim = 5*24*60*60

lbh = -2
ubh = -0.1
lbu = -1e-06  #  unit [m/s]
ubu = 0  #  unit [m/s]
# ubu = -0.32e-05  #  unit [m/s]
soilPars = loamySoil()
# lb_zone = vg.volumetricMoisture(lbh, soilPars)
# ub_zone = vg.volumetricMoisture(ubh, soilPars)
lb_zone = np.zeros(int(Nt/Delta))
ub_zone = np.zeros(int(Nt/Delta))
for i in range(int(Nt/Delta)):
    lb_zone[i] = 0.2 #min(0.18*exp(i/int(Nt/Delta)), 0.2)
    ub_zone[i] = 0.25 #max(0.23*exp(-i/int(Nt/Delta)), 0.21)
#  max((0.23+0.18)/2, 0.23 - i*(0.23-0.18)*0.15)
lby = vg.volumetricMoisture(-80)
uby = vg.volumetricMoisture(-0.1)
# Q = 4000
QL = 4000
QU = 4000
R = 10

# Change initial guess (e.g. random number)
ug = 0
yg = 0.2

# Change initial state (e.g. random number)
u0 = 0

h0 = -0.96 # -0.96 -0.15
y0 = vg.volumetricMoisture(h0)

x_i = np.ones(26)*h0

NN_past = DM.zeros((past, Nu+Ny))
for i in range(past):
    NN_past[i,0] =  (u0 - lbu)/(ubu - lbu)
    NN_past[i,1] =  (y0 - lby)/(uby - lby)
# NN_past = np.loadtxt("NN_in_0.txt")
# NN_past[0,0] = 0
# y0 = NN_past[0,1]
# NN_past[:,0] = (NN_past[:,0] - lbu)/(ubu-lbu)
# NN_past[:,1] = (NN_past[:,1] - lby)/(uby-lby)
# Define the original model for optimal 
Nx = 1
totalDepth, axialNodes, axialDistance, totalNodes = par.spatialVariables_1D()  #Spatial parameters
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = \
                                 par.temporalVariable_dataGen(25000*6*60)


# def ODE(timeArray, head, irrigAmount, cropCoeff, refEvap, soilPars):
#     origin = Richards1D_vectorized(head, irrigAmount, cropCoeff, refEvap, soilPars)
#     error = np.zeros(26)
#     # for i in range(26):
#     #     error[i] = np.random.choice([-1,1]) * 0.1*np.random.random_sample()*origin[i]
#     return origin + error

def ODE(head, irrigAmount, cropCoeff, refEvap):
    origin = Richards1D_casadi(head, irrigAmount, cropCoeff, refEvap)
    # error = DM.zeros(26)
    # for i in range(26):
    #     error[i] = np.random.choice([-1,1]) * 0.05*np.random.random_sample()*origin[i]
    return origin

t = MX.sym('t')
x = MX.sym('x', totalNodes)
I = MX.sym('I')
C = MX.sym('C')
rE = MX.sym('rE')

p = vertcat(I, C, rE)
ode = {'x': x, 'p': p, 'ode':ODE(x, p[0], p[1], p[2])}
opts = {'tf': samplingTimeInternal, 'regularity_check': True}
F = integrator('F', 'cvodes', ode, opts)

def sim_2h_MX(head, irrigAmount, cropCoeff, refEvap):
    for t in range(internalTimeSteps):
        r = F(x0 = head, p = vertcat(irrigAmount, cropCoeff, refEvap))
        head = r["xf"]
        # y = vg.volumetricMoisture(x_i[y_index])
    error = DM.zeros(26)
    # for i in range(26):
        # error[i] = np.random.choice([-1,1]) * 0.05*np.random.random_sample()*head[i]
    return head+error

cropCoeff = 0.88
refEvap = 3.102e-08

#%%  Import the weather condition 
# data = pd.read_csv('C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D\Weather.csv')
# refEvap = (data.iloc[10:20,5].values)/1e3/(24*60*60)
random.seed(5)
np.random.seed(5)
refEvap = np.random.normal(3, 0.2, 10)/1e3/(24*60*60)
Evap = np.zeros(int(Nsim/Delta))
for i in range(10):
    Evap[i*12:i*12+12] = refEvap[i]
# np.savetxt("C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D/ZMPC Simulation Results/Evap.txt", Evap)
# data_hr = pd.read_csv('C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D\Rain_hourly.csv')
# precip = data_hr.iloc[240:480,2].values/1e3/(60*60*24)

Precip = np.zeros((5,12))

for i in range(5):
    start = random.randrange(5)
    end = random.randrange(10)
    for j in range(start, end):
        Precip[i,j] = -(random.random()*(20-5)+5)/1e3/(24*60*60)
        # Precip_p[i,j] = Precip[i,j]  + np.random.choice([-1,1]) * 0.3*np.random.random_sample()*Precip[i,j]

Precip = Precip.flatten()
Precip_p = np.zeros(Precip.shape[0])
for i in range(Precip.shape[0]):
        Precip_p[i] = min(0, Precip[i]  + np.random.choice([-1,1])* 0.2*np.random.random_sample()*min(Precip))

Precip_p = np.concatenate((Precip_p.flatten(), np.zeros(20)))
# Precip_p = -abs((Precip*np.random.normal(1,0.8,size = (Precip.shape)))) # uniform

# Precip = np.zeros(int(Nsim/Delta))

# for i in range(int(Nsim/Delta)):
#     Precip[i] = random.choice([0,0,0,0,1]) * (random.random()*(2-0.2)+0.2)/1e3/(24*60*60)

plt.figure(3)
plt.step(np.arange(Precip.shape[0]), Precip)
plt.step(np.arange(Precip.shape[0]), Precip_p[:60])
# np.savetxt("C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D/ZMPC Simulation Results/Precip.txt", Precip)
# lzone_store = DM.zeros((int(Nsim/Delta),Ny))
# uzone_store = DM.zeros((int(Nsim/Delta),Ny))
# start = time.time()

# %%

start = time.time()
u = SX.sym('u',(int(Nt/Delta), Nu))
y = SX.sym('y',(int(Nt/Delta)+1, Ny))
# yz = SX.sym('yz',(int(Nt/Delta), Ny))
para = SX.sym('para',(int(Nt/Delta), Np)) 
yl = SX.sym('yl',(int(Nt/Delta), Ny))
yu = SX.sym('yu',(int(Nt/Delta), Ny))
# sUB = SX.sym('sUB',(int(Nt/Delta), Ny))

u_opt = DM.zeros((int(Nsim/Delta),Nu))
y_opt = DM.zeros((int(Nsim/Delta)+1,Ny))
# y_z = DM.zeros((int(Nsim/Delta),Ny))
y_u = DM.zeros((int(Nsim/Delta),Ny))
y_l = DM.zeros((int(Nsim/Delta),Ny))
y_opt[0,:] = y0
y_pred = DM.zeros((int(Nsim/Delta)+1,Ny))
y_pred[0,:] = y0
y_e = DM.zeros((int(Nsim/Delta)+2,Ny))
# y_e[0:2] = 0.05
# add = DM.zeros((int(Nsim/Delta),Ny))
# e_as = [] 
b_save = []
# e_a = 1
b = 0
error = []
# u, y, yz, ug, yg, lbu, ubu, lby, uby, lbZone, ubZone, 
            # Q, R, NN_past, y0, Nt, Delta, e_a, e_b
update = 2

for i in range(int(Nsim/Delta)): # int(Nsim/Delta)
    # add[i] = (y_e[i]+y_e[i+1])/2
    e = 0
    W, P, W0, W_lb, W_ub, f, g, g_lb, g_ub = zoneMPC(u, para, y, yl, yu, ug, yg, 
                                                  lbu, ubu, lby, uby, lb_zone, ub_zone, 
                                                  QL, QU, R, NN_past, y0, Nt, Delta, 1, b)
    
    nlp = {'x':W, 'p': P, 'f':f, 'g':g} # P
    S = nlpsol('S','ipopt',nlp)
    # Pi  = Precip_p[i:i+int(Nt/Delta)].reshape((int(Nt/Delta),1))
    Pi = np.zeros((int(Nt/Delta),1))
    r = S(x0 = W0, lbx = W_lb, ubx = W_ub, lbg = g_lb, ubg = g_ub, p = Pi)   
    u_i = r['x'][Ny:Nu+Ny]*(ubu - lbu) + lbu
    y_pred[i+1] = r['x'][Nu+Ny:Nu+Ny*2]
    y_l[i] = r['x'][Nu+Ny*2:Nu+Ny*3]
    y_u[i] = r['x'][Nu+Ny*3:Nu+Ny*4]
    u_opt[i,:] = u_i  #*(ubu - lbu) + lbu
    # Try NN here - no model mismatch 
    # y_i = r['x'][Nu+Ny:Nu+2*Ny]
    # NN_t = []
    # NN_t = vertcat(NN_t, M1(NN_update))
    # NN_t = vertcat(NN_t, M2(NN_update))
    # NN_t = vertcat(NN_t, M3(NN_update))
    # y_i = Ma(NN_t)*(uby - lby) + lby
    
    x_i = sim_2h_MX(x_i, u_i, cropCoeff, refEvap) # Evap[i]   + Precip[i]
    y_i = vg.volumetricMoisture(x_i[y_index]) #+ np.random.choice([-1,1]) * 0.01*((uby-lby)*np.random.random_sample() + lby)
    # for j in range(int(Delta/samplingTimeInternal)):
    #     sol = integrate.solve_ivp(ODE, args = (u_i, cropCoeff, refEvap, soilPars),
    #                           t_span=[0,samplingTimeInternal],y0=tuple(x_i),method='RK45') # + Precip[i]
    
    #     x_i = sol.y[:,-1]
        # u_i = 0  # Precip[i]
    
    # y_i = vg.volumetricMoisture(x_i[y_index], soilPars) #
    y0 = y_i
    y_e[i+2] = y_i - y_pred[i+1]
    
    
    
    NN_update = DM.zeros((1,2))
    NN_update[0,0] = r['x'][Ny:Nu+Ny]   # (u_i - lbu)/(ubu - lbu)
    NN_update[0,1] = (y_i - lby)/(uby - lby)  # (y0 - lby)/(uby - lby)
    NN_past = vertcat(NN_past[1:,:], NN_update)
    
    y_opt[i+1,:] = y_i
    e += y_i - y_pred[i+1]
# # error.append(e)
# # e = 0

    if (i+2)%update == 0:
        # b = e/update  
        b = sign(e)*min(abs(e/update), 0.2)
        b_save.append(b)
        error.append(e)
        e = 0
    # if (i+1)%12 ==0:
        # e_a, e_b = error_model(np.array(y_opt[i-11:i+1]), np.array(y_pred[i-11:i+1]))
    # e_as.append(e_a)
    # e_bs.append(e_b)
    
print((time.time()-start)/int(Nsim/Delta))

# u_sol = r['x'][1]*(ubu - lbu) + lbu
# y_sol = r['x'][2]
# for i in range(1,int(Nt/Delta)):
#     u_sol = vertcat(u_sol, r['x'][1+i*4]*(ubu - lbu) + lbu)
#     y_sol = vertcat(y_sol, r['x'][2+i*4])
# y_sol = vertcat(y_sol,r['x'][6])
# y_sol = vertcat(y_sol,r['x'][10])
# y_sol = vertcat(y_sol,r['x'][14])
# y_sol = vertcat(y_sol,r['x'][18])

# u_sol = r['x'][1]
# u_sol = vertcat(u_sol,r['x'][5])
# u_sol = vertcat(u_sol,r['x'][9])
# u_sol = vertcat(u_sol,r['x'][13])
# u_sol = vertcat(u_sol,r['x'][17])
#  
obj_opt = 0
for i in range(int(Nsim/Delta)):
    
    uc = (u_opt[i] - lbu)/(ubu-lbu)
    obj_opt += obj(uc, y_l[i], y_u[i], QL, QU, R)
    
       # 507.027 2-layered
    # obj(u_opt[i], DM(max(y_opt[i] - lb_zone[i],0)), DM(max(ub_zone[i]-y_opt[i],0)), QL, QU,R)

csfont = {'fontname':'Times New Roman'}
font = font_manager.FontProperties(family='Times New Roman',
                                    # weight='bold',
                                    style='normal', size=15)


tplot = np.linspace(0, Nsim/3600, num = int(Nsim/Delta), endpoint=False)

# weather + 12 st updates 12.401
# more rain + 6 st updates 22.987
# more rain + 12 st updates 10.0228
# shrinking zone - 0.15 - 134.372
a = 30
plt.figure(3)
# plt.plot(tplot, y_opt[:-1], label = "Exponential Shrinking - 2")
plt.plot(tplot, y_opt[:-1], label = "LSTM")
# plt.plot(tplot, y2[:-1], label = "2-layered")
plt.legend(frameon=False, ncol=1)
plt.xlabel('Time Steps', **csfont, fontsize = 17)
plt.ylabel('Soil Moisture ($m^3/m^3$)', **csfont, fontsize = 17)


plt.figure(4)
plt.step(tplot, -u_opt , label = "LSTM")#, label = "Exponential Shrinking - 2"
# plt.step(tplot, -u_2L, label = "Sym. Zone")
plt.legend(frameon=False, ncol=1)
plt.xlabel('Time Steps', **csfont, fontsize = 17)
plt.ylabel('Irrigation (m/s)', **csfont, fontsize = 17)

# plt.figure(1)
# plt.plot(lb_zone, "b-.", label = "Exponential 2")
# plt.plot(ub_zone, "b-.")


# plt.figure(6)
# plt.plot(y_e[1:], label = "Rain + 10% P Noise + 2% M Noise - no correction")

# np.savetxt("C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D/ZMPC Simulation Results/Noiminal_u.txt", u_opt)
# np.savetxt("C:/Users/zhiyinan/OneDrive - ualberta.ca/Irrigation/crop_model_1D/ZMPC Simulation Results/Noiminal_y.txt", y_opt)
# weather condition + const bias 83.462
# update every 6 sp, 2 coeff 20.7828
# 
# add = np.zeros(int(Nsim/Delta))
# for i in range(int(Nsim/Delta)):
#     add[i] =  (y_e[i]+y_e[i+1])/2
# %% 

# NN_past = np.loadtxt("NN_in_0.txt")
# NN_past[0,0] = 0
# y0 = NN_past[0,1]
# NN_past[:,0] = (NN_past[:,0] - lbu)/(ubu-lbu)
# NN_past[:,1] = (NN_past[:,1] - lby)/(uby-lby)
# yp = np.zeros(11)
# yp[0] = y0
# NI = NN_past
# for i in range(10):
#     uc = (u_opt[i] - lbu)/(ubu-lbu)
#     yc = (yp[i] - lby)/(uby-lby)
#     NI = vertcat(NI[1:,:], horzcat(uc, yc))
#     yp[i+1] = M(NI) * (uby-lby) + lby

# def model(x):
#     n1 = 5
#     n2 = 3
#     # Initialise the model with the hidden states
#     h_t1 = SX.zeros((1, n1))
#     h_t2 = SX.zeros((1, n2))
    
#     # Long term memory
#     C_t1 = SX.zeros((1, n1))
#     C_t2 = SX.zeros((1, n2))
#     for i in range(x.shape[0]):
#         xi = x[i,:].reshape((1,2))
#         L1 = LSTMlayer_sig(warr1, uarr1, barr1, xi, h_t1[i,:], C_t1[i,:])
#         h_t1 = vertcat(h_t1, L1[0])
#         C_t1 = vertcat(C_t1, L1[1])
    
#     for i in range(x.shape[0]):
#         L2 = LSTMlayer_tanh(warr2, uarr2, barr2, h_t1[i+1,:], h_t2[i,:], C_t2[i,:])
#         h_t2 = vertcat(h_t2, L2[0])
#         C_t2 = vertcat(C_t2, L2[1])
#     #  array([-2.7764497, -2.6578658, -2.6829786, -2.7356846, -2.73696  ]
#     # Dense layer
#     out3 = []
#     for i in range(bias3.shape[0]):
#         temp3 = mtimes(h_t2[-1,:], kernel3)[i]
#     temp3 = bias3.reshape((1, bias3.shape[0])) + mtimes(h_t2[-1,:], kernel3)
#     out3 = sigmoid(temp3)     
    
#     return out3

# def model_large_sampling_time(tn,x0):
#     output = DM.zeros(tn, Ny)
#     u = x0[-1,:Nu]
#     for i in range(tn):
#         y = model(x0)
#         output[i,:] = y
#         x0 = vertcat(x0[1:,:], horzcat(u,y))
    
#     return output





